#include "optimizer_adapter/from_pg_query.h"

#include <cmath>
#include <memory>
#include <utility>
#include <vector>

#include "planner/expression.h"
#include "planner/logical_operators.h"
#include "planner/optimizer_defs.h"

extern "C" {
#include <postgres.h>

#include <nodes/parsenodes.h>
#include <nodes/pg_list.h>
#include <optimizer/clauses.h>
#include <optimizer/optimizer.h>
#include <optimizer/prep.h>
#include <parser/parsetree.h>
#include <storage/lockdefs.h>
}

namespace optimizer {

#define TODO(msg) throw std::runtime_error("Not implemented yet function: " #msg)

Query* QueryToOperatorTransformer::PreProcessQuery(Query* query) {
  transform_MERGE_to_join(query);
  replace_empty_jointree(query);
  return query;
}

std::shared_ptr<Operator> QueryToOperatorTransformer::TransformQuery(Query* query) {
  std::shared_ptr<Operator> root;

  // query = PreProcessQuery(query);

  switch (query->commandType) {
    case CMD_SELECT:
      root = TransformSelectQuery(query);
      break;

    case CMD_INSERT:
    case CMD_DELETE:
    case CMD_UPDATE:
    default:
      throw std::runtime_error("Unsupported node type");
  }

  return root;
}

/*
0. cte
1. set
2. from
  1. table or subquery or join
  2. condition
3. aggarate
  * group by
  * agg in target list
4. having
5. projection
6. order by
7. limit

* pg101: target list 指的是当前 query 需要的列，而不是最终的输出，例如 sort，group
等，运算中需要，但是实际上可以不输出， 使用 resjunk 判断
*/
std::shared_ptr<Operator> QueryToOperatorTransformer::TransformSelectQuery(Query* query) {
  if (query->targetList == nullptr)
    throw std::runtime_error("SELECT query has no target list");

  std::shared_ptr<Operator> root;
  processing_queries_.push_back(query);
  if (query->setOperations != nullptr) {
    root = VisitNode(query->setOperations);
  } else {
    root = Visit(query->jointree);
  }

  recorder_agg_exprs_ = true;
  List* sort_grouo_exprs = nullptr;
  List* target_exprs = nullptr;
  List* agg_exprs = nullptr;
  foreach_node(TargetEntry, target, query->targetList) {
    auto* expr = (target->expr);

    /* get sort and group exprs */
    if (target->ressortgroupref != 0)
      sort_grouo_exprs = lappend(sort_grouo_exprs, expr);
    if (contain_agg_clause((Node*)expr))
      agg_exprs = lappend(agg_exprs, expr);
    target_exprs = lappend(target_exprs, expr);
  }
  recorder_agg_exprs_ = false;

  /*
  pg101: 可以只有 group 而没有聚合，且允许出现其他列， 但是限定 group by 主键，相当于一个空操作，因为主键唯一
  tips: 可以优化掉
   */
  if (query->groupClause != nullptr || query->hasAggs) {
    List* group_exprs = get_sortgrouplist_exprs(query->groupClause, query->targetList);

    auto agg = std::make_shared<LogicalAggregate>(group_exprs, agg_exprs);
    if (query->havingQual != nullptr)
      agg->SetFilter(list_make1(query->havingQual));
    agg->AddChild(root);
    root = std::move(agg);
  }

  if (query->distinctClause != nullptr) {
    TODO("DISTINCT");
  }

  if (query->limitCount != nullptr) {
    auto limit = std::make_shared<LogicalLimit>((Expression*)query->limitCount, (Expression*)query->limitOffset);
    limit->AddChild(root);
    root = std::move(limit);
  }

  root->SetProject(target_exprs);
  processing_queries_.pop_back();
  return root;
}

std::shared_ptr<Operator> QueryToOperatorTransformer::VisitNode(Node* node) {
  switch (node->type) {
    case T_SetOperationStmt:
      return Visit(reinterpret_cast<SetOperationStmt*>(node));
      break;

    case T_RangeTblRef:
      return Visit(reinterpret_cast<RangeTblRef*>(node));
      break;

    case T_FromExpr:
      return Visit(reinterpret_cast<FromExpr*>(node));
      break;

    case T_JoinExpr:
      return Visit(reinterpret_cast<JoinExpr*>(node));
      break;

    default:
      throw std::runtime_error("Unsupported node type");
  }
}

/*
  JoinExpr 里面应该只有简单的常规类型的join，semi join anti join 等非常规join只能在后续 unnest 之后才会有
*/
std::shared_ptr<Operator> QueryToOperatorTransformer::Visit(JoinExpr* join_expr) {
  auto left = VisitNode(join_expr->larg);
  auto right = VisitNode(join_expr->rarg);

  auto join = std::make_shared<LogicalJoin>(join_expr->jointype, list_make1(join_expr->quals));

  join->AddChild(left);
  join->AddChild(right);
  return join;
}

std::shared_ptr<Operator> QueryToOperatorTransformer::Visit(FromExpr* from) {
  std::shared_ptr<Operator> root;
  if (from->fromlist == nullptr) {
    root = std::make_shared<LogicalGet>(nullptr);
  } else {
    // single table or join
    if (list_length(from->fromlist) == 1)
      root = VisitNode(reinterpret_cast<Node*>(linitial(from->fromlist)));
    else {
      // mutiple tables, join them
      auto ref_1 = VisitNode(reinterpret_cast<Node*>(linitial(from->fromlist)));
      for (auto i = 1; i < list_length(from->fromlist); i++) {
        auto ref_2 = VisitNode(reinterpret_cast<Node*>(list_nth(from->fromlist, i)));
        auto join = std::make_shared<LogicalJoin>(JOIN_INNER, nullptr);
        join->AddChild(ref_1);
        join->AddChild(ref_2);
        ref_1 = std::move(join);
      }
      root = std::move(ref_1);
    }
  }

  if (from->quals != nullptr)
    root->SetFilter(list_make1(from->quals));

  return root;
}
std::shared_ptr<Operator> QueryToOperatorTransformer::Visit(RangeTblRef* table_ref) {
  const auto* query = processing_queries_.back();
  auto* rte = list_nth_node(RangeTblEntry, query->rtable, table_ref->rtindex - 1);
  switch (rte->rtekind) {
    case RTE_RELATION:
      return std::make_shared<LogicalGet>(rte);
    case RTE_SUBQUERY:
      return TransformSelectQuery(rte->subquery);
    case RTE_JOIN:
    case RTE_FUNCTION:
    case RTE_TABLEFUNC:
    case RTE_VALUES:
    case RTE_CTE:
    case RTE_NAMEDTUPLESTORE:
    case RTE_RESULT:

    default:
      throw std::runtime_error("Unsupported RTE type");
  }
  return nullptr;
}

std::shared_ptr<Operator> QueryToOperatorTransformer::Visit(SetOperationStmt* set) {
  auto left = VisitNode(set->larg);
  auto right = VisitNode(set->rarg);

  OperatorType op_type = OperatorType::UNDEFINED;

  switch (set->op) {
    case SETOP_UNION:
      op_type = OperatorType::LogicalUnion;
      break;
    case SETOP_INTERSECT:
      op_type = OperatorType::LogicalIntersect;
      break;
    case SETOP_EXCEPT:
      op_type = OperatorType::Logicalexcept;
      break;
    case SETOP_NONE:
      throw std::runtime_error("Unsupported set operation type");
  }

  auto logical_set = std::make_shared<LogicalSetOperation>(op_type, set->all);

  auto root = std::move(logical_set);
  root->AddChild(left);
  root->AddChild(right);

  return root;
}

}  // namespace optimizer
